{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c9308b7",
   "metadata": {},
   "source": [
    "# 测试自定义算子合法化变换功能\n",
    "\n",
    "本测试验证用户是否可以为特定算子定义自定义的合法化变换逻辑。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1a85f150",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.relax.transform import LegalizeOps\n",
    "from tvm.relax.transform.legalize_ops.common import register_legalize\n",
    "from tvm.script import relax as R, tir as T, ir as I"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e0e989",
   "metadata": {},
   "source": [
    "## 自定义 legalize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed7bf734",
   "metadata": {},
   "source": [
    "将 `R.add` 算子通过自定义函数变换为使用 `topi.add` 的 TE 调用并验证变换后的 IR 是否与预期结果一致。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "65a193b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Add:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((1, 2, 3), \"float32\"), y: R.Tensor((4, 3, 2, 1), \"float32\")) -> R.Tensor((4, 3, 2, 3), \"float32\"):\n",
    "        gv: R.Tensor((4, 3, 2, 3), \"float32\") = R.add(x, y)\n",
    "        return gv\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fa2cf92",
   "metadata": {},
   "source": [
    "定义自定义的合法化变换函数，将 `x` 和 `y` 的顺序调换后调用 `topi.add`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9f0c3b48",
   "metadata": {},
   "outputs": [],
   "source": [
    "def customize_legalize_add(bb: relax.BlockBuilder, call: relax.Call):\n",
    "    from tvm import topi  # pylint: disable=import-outside-toplevel\n",
    "    return bb.call_te(topi.add, call.args[1], call.args[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "618cf5ca",
   "metadata": {},
   "source": [
    "应用 `LegalizeOps` 变换，传入自定义的合法化函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "791e5514",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">add</span>(y: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), x: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_add: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1, ax2, ax3 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">4</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">2</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1, v_ax2, v_ax3 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSSS&quot;</span>, [ax0, ax1, ax2, ax3])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(y[v_ax0, v_ax1, v_ax2, T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">0</span>)], x[T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">0</span>), v_ax2, v_ax3])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])\n",
       "                T_add[v_ax0, v_ax1, v_ax2, v_ax3] <span style=\"color: #A2F; font-weight: bold\">=</span> y[v_ax0, v_ax1, v_ax2, T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">0</span>)] <span style=\"color: #A2F; font-weight: bold\">+</span> x[T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">0</span>), v_ax2, v_ax3]\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>add, (y, x), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod = LegalizeOps({\"relax.add\": customize_legalize_add})(Add)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c836642d",
   "metadata": {},
   "source": [
    "## 测试不同类型的调用是否都能被正确合法化"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84b82181",
   "metadata": {},
   "source": [
    "定义包含多种调用类型的测试模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "689bf490",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def mul2(x: R.Tensor((3, 3), \"float32\")):\n",
    "        gv = R.multiply(x, R.const(2.0, \"float32\"))\n",
    "        return gv\n",
    "\n",
    "    @T.prim_func(private=True)\n",
    "    def identity(rxplaceholder: T.Buffer((T.int64(3), T.int64(3)), \"float32\"), T_id: T.Buffer((T.int64(3), T.int64(3)), \"float32\")):\n",
    "        for ax0, ax1 in T.grid(T.int64(3), T.int64(3)):\n",
    "            with T.block(\"T_add\"):\n",
    "                v_ax0, v_ax1 = T.axis.remap(\"SS\", [ax0, ax1])\n",
    "                T.reads(rxplaceholder[v_ax0, v_ax1])\n",
    "                T.writes(T_id[v_ax0, v_ax1])\n",
    "                T_id[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1]\n",
    "\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 3), \"float32\")):\n",
    "        cls = Before\n",
    "        gv: R.Tensor((3, 3), \"float32\") = cls.mul2(x)\n",
    "        gv1 = R.call_tir(cls.identity, gv, R.Tensor((3, 3), dtype=\"float32\"))\n",
    "        gv2 = R.multiply(gv1, R.const(2.0, \"float32\"))\n",
    "        return gv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fd285f64",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">identity</span>(rxplaceholder: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_id: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(rxplaceholder[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_id[v_ax0, v_ax1])\n",
       "                T_id[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> rxplaceholder[v_ax0, v_ax1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">multiply</span>(gv1: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_multiply: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_multiply&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(gv1[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_multiply[v_ax0, v_ax1])\n",
       "                T_multiply[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> gv1[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2.0</span>)\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">mul2</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>multiply, (x,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> cls<span style=\"color: #A2F; font-weight: bold\">.</span>mul2(x)\n",
       "        gv1 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>identity, (gv,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        gv2 <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>multiply, (gv1,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv2\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 应用LegalizeOps转换\n",
    "After = LegalizeOps()(Before)\n",
    "After.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e23d62f",
   "metadata": {},
   "source": [
    "## 测试无法进行合法化转变换的情况"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c97128c3",
   "metadata": {},
   "source": [
    "本测试验证当算子没有对应的合法化函数或缺少必要的形状信息时，变换行为是否符合预期"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "013a66f4",
   "metadata": {},
   "source": [
    "情况1：算子没有对应的合法化函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e236e3c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "add_legalize = tvm.ir.Op.get(\"relax.add\").get_attr(\"FLegalize\")\n",
    "# 重置属性用于测试\n",
    "tvm.ir.Op.get(\"relax.add\").reset_attr(\"FLegalize\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d46ce04a",
   "metadata": {},
   "source": [
    "定义简单的包含 `add` 算子的模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e23e654b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class Before0:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 3), \"float32\")):\n",
    "        gv: R.Tensor((3, 3), \"float32\") = R.add(x, x)\n",
    "        return gv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b0f187bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 应用转换（此时没有add操作的合法化函数）\n",
    "After0 = LegalizeOps()(Before0)\n",
    "# 验证模块是否保持不变\n",
    "tvm.ir.assert_structural_equal(After0, Before0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6acf9939",
   "metadata": {},
   "source": [
    "恢复原有的合法化函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e033611d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ffi.Function(0x55de62489d30)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "register_legalize(\"relax.add\", add_legalize)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05013cd1",
   "metadata": {},
   "source": [
    "情况2：无法确定所有形状信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "63011a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = relax.Var(\"s\", relax.ShapeStructInfo((3, 3)))\n",
    "x = relax.Var(\"x\", relax.TensorStructInfo((3, 3), \"float32\"))\n",
    "y = relax.Var(\"y\", relax.TensorStructInfo(s, \"float32\"))\n",
    "bb = relax.BlockBuilder()\n",
    "with bb.function(\"main\", [x, y]):\n",
    "    with bb.dataflow():\n",
    "        gv = bb.emit_output(R.add(x, y))\n",
    "    bb.emit_func_output(gv)\n",
    "Before1 = bb.get()\n",
    "# 应用转换（此时无法确定y的完整形状信息）\n",
    "After1 = LegalizeOps()(Before1)\n",
    "# 验证模块是否保持不变\n",
    "tvm.ir.assert_structural_equal(After1, Before1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc9e37e2",
   "metadata": {},
   "source": [
    "## 测试不同数据类型的标量算子合法化变换"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6901f5a0",
   "metadata": {},
   "source": [
    "本测试验证 `LegalizeOps` 在处理 `float16`、`uint8` 和 `bool` 等不同数据类型时能够正确保留类型信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "349cd448",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义float16数据类型的测试模块\n",
    "@tvm.script.ir_module\n",
    "class Before0:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 3), \"float16\")):\n",
    "        gv: R.Tensor((3, 3), \"float16\") = R.multiply(x, R.const(1.14514, \"float16\"))\n",
    "        return gv\n",
    "\n",
    "# 定义uint8数据类型的测试模块\n",
    "@tvm.script.ir_module\n",
    "class Before1:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 3), \"uint8\")):\n",
    "        gv: R.Tensor((3, 3), \"uint8\") = R.multiply(x, R.const(2, \"uint8\"))\n",
    "        return gv\n",
    "\n",
    "# 定义bool数据类型的测试模块\n",
    "@tvm.script.ir_module\n",
    "class Before2:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 3), \"bool\")):\n",
    "        gv: R.Tensor((3, 3), \"bool\") = R.equal(x, R.const(True, \"bool\"))\n",
    "        return gv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7797fd4",
   "metadata": {},
   "source": [
    "应用转换并验证结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "75ea75e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "After0 = LegalizeOps()(Before0)\n",
    "After1 = LegalizeOps()(Before1)\n",
    "After2 = LegalizeOps()(Before2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "97ba5ccb",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">multiply</span>(x: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float16&quot;</span>), T_multiply: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;float16&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_multiply&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(x[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_multiply[v_ax0, v_ax1])\n",
       "                T_multiply[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> x[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>float16(<span style=\"color: #008000\">1.1455078125</span>)\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>multiply, (x,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float16&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "After0.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "af04e78c",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">multiply</span>(x: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;uint8&quot;</span>), T_multiply: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;uint8&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_multiply&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(x[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_multiply[v_ax0, v_ax1])\n",
       "                T_multiply[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> x[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">*</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>uint8(<span style=\"color: #008000\">2</span>)\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;uint8&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;uint8&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>multiply, (x,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;uint8&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "After1.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "20eb557b",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">equal</span>(x: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;bool&quot;</span>), T_equal: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)), <span style=\"color: #BA2121\">&quot;bool&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">3</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_equal&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(x[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_equal[v_ax0, v_ax1])\n",
       "                T_equal[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> x[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">==</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;bool&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;bool&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>equal, (x,), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;bool&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "After2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82fec25",
   "metadata": {},
   "source": [
    "## 测试矩阵乘法算子合法化要求已知数据类型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "200b1e7e",
   "metadata": {},
   "source": [
    "本测试验证当 `matmul` 算子缺少明确数据类型时，合法化转换应抛出适当的错误"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "92cb6aeb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[13:44:41] /media/pc/data/lxw/ai/tvm/src/relax/ir/block_builder.cc:64: Warning: BlockBuilder destroyed with remaining blocks!\n"
     ]
    }
   ],
   "source": [
    "import pytest\n",
    "@I.ir_module\n",
    "class ArbitraryDtype:\n",
    "    @R.function\n",
    "    def main(A: R.Tensor([16, 32]), B: R.Tensor([32, 8])) -> R.Tensor([16, 8]):\n",
    "        return R.matmul(A, B)\n",
    "\n",
    "# 验证转换时是否抛出预期的错误\n",
    "with pytest.raises(AssertionError) as err:\n",
    "    LegalizeOps()(ArbitraryDtype)\n",
    "\n",
    "# 错误应该是在尝试合法化R.matmul时捕获的，并提供友好的错误消息\n",
    "# 而不是等到`BlockBuilder.call_te`实现时，尝试创建kHandle类型的数值常量时才抛出错误\n",
    "err_message = err.value.args[0]\n",
    "assert err_message.startswith(\"To legalize R.matmul\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b3543c4",
   "metadata": {},
   "source": [
    "## 测试带 vdevice 的合法化变换\n",
    "\n",
    "本测试验证当参数类型仅在 `vdevice` 上不同时，`LegalizeOps` 能够为不同目标生成不同的内核。 \n",
    "\n",
    "这是回归测试，之前的实现中，具有不同 `vdevice` 的参数类型会被合法化为使用相同的 `PrimFunc`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "df0e1461",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    I<span style=\"color: #A2F; font-weight: bold\">.</span>module_global_infos({<span style=\"color: #BA2121\">&quot;vdevice&quot;</span>: [I<span style=\"color: #A2F; font-weight: bold\">.</span>vdevice({<span style=\"color: #BA2121\">&quot;keys&quot;</span>: [<span style=\"color: #BA2121\">&quot;cpu&quot;</span>], <span style=\"color: #BA2121\">&quot;kind&quot;</span>: <span style=\"color: #BA2121\">&quot;llvm&quot;</span>, <span style=\"color: #BA2121\">&quot;mtriple&quot;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>, <span style=\"color: #BA2121\">&quot;tag&quot;</span>: <span style=\"color: #BA2121\">&quot;&quot;</span>}, <span style=\"color: #008000\">0</span>, <span style=\"color: #BA2121\">&quot;global&quot;</span>)]})\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">add</span>(A: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_add: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> A[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">+</span> B[v_ax0, v_ax1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@T</span><span style=\"color: #A2F; font-weight: bold\">.</span>prim_func(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">add_llvm</span>(A: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>), T_add: T<span style=\"color: #A2F; font-weight: bold\">.</span>Buffer((T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        T<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;target&quot;</span>: T<span style=\"color: #A2F; font-weight: bold\">.</span>target({<span style=\"color: #BA2121\">&quot;keys&quot;</span>: [<span style=\"color: #BA2121\">&quot;cpu&quot;</span>], <span style=\"color: #BA2121\">&quot;kind&quot;</span>: <span style=\"color: #BA2121\">&quot;llvm&quot;</span>, <span style=\"color: #BA2121\">&quot;mtriple&quot;</span>: <span style=\"color: #BA2121\">&quot;x86_64-unknown-linux-gnu&quot;</span>, <span style=\"color: #BA2121\">&quot;tag&quot;</span>: <span style=\"color: #BA2121\">&quot;&quot;</span>}), <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>grid(T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>), T<span style=\"color: #A2F; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">32</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>axis<span style=\"color: #A2F; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])\n",
       "                T<span style=\"color: #A2F; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">=</span> A[v_ax0, v_ax1] <span style=\"color: #A2F; font-weight: bold\">+</span> B[v_ax0, v_ax1]\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">func_cuda</span>(A: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), B: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        C <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>add, (A, B), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> C\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">func_llvm</span>(A: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>, vdevice<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;llvm:0&quot;</span>), B: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>, vdevice<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;llvm:0&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>, vdevice<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;llvm:0&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        C <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>call_tir(cls<span style=\"color: #A2F; font-weight: bold\">.</span>add_llvm, (A, B), out_sinfo<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>, vdevice<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;llvm:0&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> C\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@I.ir_module\n",
    "class Before:\n",
    "    I.module_global_infos({\"vdevice\": [I.vdevice(\"llvm\")]})\n",
    "\n",
    "    @R.function\n",
    "    def func_cuda(A: R.Tensor([32, 32], \"float32\"), B: R.Tensor([32, 32], \"float32\")):\n",
    "        C = R.add(A, B)\n",
    "        return C\n",
    "\n",
    "    @R.function\n",
    "    def func_llvm(\n",
    "        A: R.Tensor([32, 32], \"float32\", \"llvm\"), B: R.Tensor([32, 32], \"float32\", \"llvm\")\n",
    "    ):\n",
    "        C = R.add(A, B)\n",
    "        return C\n",
    "\n",
    "# 在CUDA目标下应用转换\n",
    "with tvm.target.Target(\"cuda\"):\n",
    "    After = tvm.relax.transform.LegalizeOps()(Before)\n",
    "After.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcfd53a4",
   "metadata": {},
   "source": [
    "## 自定义算子，测试递归合法化功能"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34ada24b",
   "metadata": {},
   "source": [
    "定义测试参数，测试不同的合法化变换返回方式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f6415d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm.testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ee67e5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def register_custom_op(emit_legalization_through_builder):\n",
    "    op_name = \"custom_op.matmul_bias_add\"\n",
    "\n",
    "    # 定义结构信息推断函数\n",
    "    def infer_struct_info(call: relax.Call, context):\n",
    "        activations, weight, bias = call.args\n",
    "\n",
    "        matmul_call = relax.op.matmul(activations, weight)\n",
    "        matmul_sinfo = tvm.ir.Op.get(\"relax.matmul\").get_attr(\"FInferStructInfo\")(\n",
    "            matmul_call, context\n",
    "        )\n",
    "\n",
    "        matmul_var = relax.Var(\"dummy_var\", matmul_sinfo)\n",
    "        add_call = matmul_var + bias\n",
    "        add_sinfo = tvm.ir.Op.get(\"relax.add\").get_attr(\"FInferStructInfo\")(add_call, context)\n",
    "\n",
    "        return add_sinfo\n",
    "\n",
    "    # 定义合法化函数\n",
    "    def legalize(bb: relax.BlockBuilder, call: relax.Call):\n",
    "        activations, weight, bias = call.args\n",
    "        legalized = relax.op.matmul(activations, weight) + bias\n",
    "        if emit_legalization_through_builder:\n",
    "            legalized = bb.emit(legalized)\n",
    "        return legalized\n",
    "\n",
    "    # 注册操作的属性\n",
    "    op_attrs = {\n",
    "        \"FInferStructInfo\": infer_struct_info,\n",
    "        \"FLegalize\": legalize,\n",
    "        \"FPurity\": True,\n",
    "    }\n",
    "\n",
    "    for key, value in op_attrs.items():\n",
    "        tvm.ir.register_op_attr(op_name, key, value)\n",
    "\n",
    "    op = tvm.ir.Op.get(op_name)\n",
    "    return op\n",
    "    # yield op\n",
    "\n",
    "    # # 清理：重置属性\n",
    "    # for key in op_attrs:\n",
    "    #     op.reset_attr(key)\n",
    "\n",
    "\n",
    "custom_op = register_custom_op(emit_legalization_through_builder={\n",
    "    \"return_relax_expr\": False,\n",
    "    \"return_relax_var\": True,\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf5851a9",
   "metadata": {},
   "source": [
    "本测试验证算子的合法化变换可能生成新的需要合法化的算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "519fe8dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        A: R.Tensor([16, 32, 64], \"float32\"),\n",
    "        Weight: R.Tensor([64, 128], \"float32\"),\n",
    "        Bias: R.Tensor([16, 32, 128], \"float32\"),\n",
    "    ):\n",
    "        return relax.Call(custom_op, [A, Weight, Bias])\n",
    "\n",
    "# 应用一次LegalizeOps转换\n",
    "AfterFirstIter = LegalizeOps()(Before)\n",
    "# 再次应用LegalizeOps转换\n",
    "AfterSecondIter = LegalizeOps()(AfterFirstIter)\n",
    "\n",
    "# LegalizeOps后，自定义操作应被替换为`R.matmul`和`R.add`，\n",
    "# 这些操作又应被替换为TIR实现。因此，第二次应用LegalizeOps()应该是无效操作。\n",
    "tvm.ir.assert_structural_equal(AfterFirstIter, AfterSecondIter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d318d41",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
